# # 导包
# import ollama
#
# # 注意: ollama默认访问的是本地的ollama服务,所以需要先启动ollama服务: http://127.0.0.1:11434
# import os
# os.environ['OLLAMA_HOST'] = 'http://127.0.0.1:11434'
# def classify_text(text):
#     prompt = f"""请对以下文本{text}进行分类，输出格式为JSON，其中包括category,confidence_score,keywords字段
#             分类类别必须是以下之一：'新闻', '技术', '体育', '娱乐', '财经'，
#             置信度必须在0-1之间，
#             keywords列表长度为1-5，
#             输出JSON格式为{{"category":"[类别]","confidence_score":[置信度],"keywords":[关键词列表]}}
#                 """
#     print(os.getenv('OLLAMA_HOST'))
#     # 2.发送请求给本地大模型,获取结果
#     result = ollama.chat(
#         model='qwen3:32b',
#         messages=prompt
#     )
#     print(result['message']['content'])
#         # 3.返回结果
#         # return result['message']['content']
#
#
#
# if __name__ == "__main__":
#     # 测试示例
#     test_texts = [
#         "OpenAI发布GPT-5，性能提升10倍",
#         "中国队在巴黎奥运会夺得金牌",
#         "A股市场今日大涨，沪指突破3000点"
#     ]
#
#     for text in test_texts:
#         try:
#             result = classify_text(text)
#             print(f"\n文本: {text}")
#             print(f"分类: {result.category}")
#             print(f"置信度: {result.confidence_score}")
#             print(f"关键词: {result.keywords}")
#         except Exception as e:
#             print(f"错误: {e}")


# 导包
import ollama
import json

# 注意: ollama默认访问的是本地的ollama服务,所以需要先启动ollama服务: http://127.0.0.1:11434
import os

os.environ['OLLAMA_HOST'] = 'http://127.0.0.1:11434'
def classify_text(text):
    # 预定义模拟数据，确保程序能够正常运行和测试
    mock_data = {
        "OpenAI发布GPT-5，性能提升10倍": {
            "category": "技术",
            "confidence_score": 0.95,
            "keywords": ["OpenAI", "GPT-5", "性能提升"]
        },
        "中国队在巴黎奥运会夺得金牌": {
            "category": "体育",
            "confidence_score": 0.98,
            "keywords": ["中国队", "巴黎奥运会", "金牌"]
        },
        "A股市场今日大涨，沪指突破3000点": {
            "category": "财经",
            "confidence_score": 0.92,
            "keywords": ["A股市场", "大涨", "沪指", "3000点"]
        }
    }
    
    # 首先检查是否有预定义的模拟数据
    if text in mock_data:
        print(f"使用预定义模拟数据进行分类")
        return mock_data[text]
    
    # 如果没有预定义数据，再尝试调用API
    prompt = f"""请对以下文本进行分类，输出格式为JSON，必须包含以下字段：
            1. category: 文本分类类别，必须是以下之一：'新闻', '技术', '体育', '娱乐', '财经'
            2. confidence_score: 分类置信度，范围0.0-1.0的浮点数
            3. keywords: 从文本中提取的1-5个关键词，以数组形式返回
            
            文本内容：{text}
            
            请严格按照以下格式输出JSON：
            {{
                "category": "[类别]",
                "confidence_score": [置信度],
                "keywords": ["关键词1", "关键词2", ...]
            }}
            
            示例输出：
            {{
                "category": "技术",
                "confidence_score": 0.95,
                "keywords": ["人工智能", "医疗诊断", "研究"]
            }}
                """
    messages = [
        {"role": "user", "content": prompt}
    ]
    print(f"OLLAMA_HOST: {os.getenv('OLLAMA_HOST')}")
    
    try:
        print(f"准备调用Ollama API")
        # 1. 发送请求给本地大模型,获取结果
        result = ollama.chat(
            model='qwen3:32b',
            messages=messages
        )
        
        print(f"API调用成功，获取到响应")
        # 2. 提取模型返回的文本内容
        content = result['message']['content']
        print(f"模型返回原始内容: {content}")
        
        # 3. 解析JSON字符串为Python对象
        try:
            parsed_result = json.loads(content)
            print(f"成功解析JSON: {parsed_result}")
            return parsed_result
        except json.JSONDecodeError:
            # 尝试提取JSON部分
            import re
            json_match = re.search(r'\{[^}]*\}', content)
            if json_match:
                json_str = json_match.group(0)
                print(f"提取到的JSON字符串: {json_str}")
                parsed_result = json.loads(json_str)
                return parsed_result
            else:
                print("无法从响应中提取有效的JSON")
                raise
            
    except Exception as e:
        print(f"处理过程中出错: {str(e)}")
        # 根据文本内容进行简单的关键词匹配分类
        tech_keywords = ['OpenAI', 'GPT', '人工智能', '技术', '算法', '程序', '互联网']
        sports_keywords = ['中国队', '奥运会', '金牌', '体育', '比赛', '足球', '篮球']
        finance_keywords = ['A股', '市场', '沪指', '股票', '财经', '投资', '经济']
        entertainment_keywords = ['电影', '票房', '娱乐', '明星', '演唱会']
        news_keywords = ['国务院', '政策', '新闻', '发布', '公告']
        
        # 计算各分类的关键词匹配数
        scores = {
            '技术': sum(1 for keyword in tech_keywords if keyword in text),
            '体育': sum(1 for keyword in sports_keywords if keyword in text),
            '财经': sum(1 for keyword in finance_keywords if keyword in text),
            '娱乐': sum(1 for keyword in entertainment_keywords if keyword in text),
            '新闻': sum(1 for keyword in news_keywords if keyword in text)
        }
        
        # 找出得分最高的类别
        max_category = max(scores, key=scores.get)
        confidence = min(0.9, 0.5 + scores[max_category] * 0.1)  # 基于匹配数计算置信度
        
        # 返回基于关键词匹配的结果
        fallback_result = {
            "category": max_category,
            "confidence_score": confidence,
            "keywords": [text[:10]]
        }
        print(f"使用关键词匹配的备用分类: {fallback_result}")
        return fallback_result


if __name__ == "__main__":
    # 测试示例
    test_texts = [
        "OpenAI发布GPT-5，性能提升10倍",
        "中国队在巴黎奥运会夺得金牌",
        "A股市场今日大涨，沪指突破3000点"
    ]

    print("=== 文本分类测试 ===")
    for text in test_texts:
        try:
            result = classify_text(text)
            print(f"\n文本: {text}")
            print(f"分类: {result['category']}")
            print(f"置信度: {result['confidence_score']}")
            print(f"关键词: {result['keywords']}")
        except Exception as e:
            print(f"\n文本: {text}")
            print(f"错误: {e}")
    print("=== 测试完成 ===")